import json
import traceback
from typing import List, Dict, Any
from .db_manager import DatabaseManager
from .ai_config import OLLAMA_CONFIG, PROMPT_TEMPLATE
from ollama import chat, Client
import os
from dotenv import load_dotenv
from .ai_query_base import AIQueryBase

# Load environment variables from .env file
load_dotenv()

class OllamaQueryManager(AIQueryBase):
    """
    OllamaQueryManager 用于通过本地 OLLAMA API 将自然语言问题转为 SQL 并查询数据库。
    负责生成提示词、调用本地 AI 接口、处理 SQL 及执行数据库查询。
    """
    def __init__(self):
        """
        初始化 OllamaQueryManager，创建数据库管理器并打印当前 OLLAMA 配置信息。
        """
        super().__init__()
        # 打印配置信息，帮助调试
        print(f"OLLAMA配置: {OLLAMA_CONFIG}")

    
    def _get_ai_response(self, prompt: str) -> str:
        """
        调用本地 OLLAMA API 获取 SQL 响应。
        参数：
            prompt (str): AI prompt。
        返回：
            str: AI 返回的 SQL 语句。
        失败时返回 None。
        """
        try:
            print(f"正在调用本地OLLAMA API，使用模型: {OLLAMA_CONFIG['model']}")
            
            # 使用ollama包的chat函数
            client = Client(
                host=OLLAMA_CONFIG['host']
            )
            response = client.chat(
                model=OLLAMA_CONFIG['model'],
                messages=[
                    {
                        "role": "assistant",
                        "content": "你是一个SQL专家，负责将自然语言转换为SQL查询语句。请只返回SQL语句，不要包含任何其他说明。\n"
                     +"请注意针对PostgreSQL数据库，请使用正确的语法。例如涉及到user等是PostgreSQL的保留关键字相关的数据库表名或字段，需要使用双引号包裹。（例如虽然是system_user，但 PostgreSQL 在某些解析情况下仍然容易对类似的名字卡壳，需要使用双引号包裹system_user）。"
                     +"请注意正确理解个字段的数据类型，避免错误数据转换。\n"
                     +"请直接给出最终答案，不要解释过程。"
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                options={
                    "temperature": OLLAMA_CONFIG['temperature']
                }
            )
            
            # 打印原始响应内容，帮助调试
            print(f"原始响应内容: {response['message']['content'][:200]}...")  # 只打印前200个字符
            
            # 提取消息内容
            content = response['message']['content'].strip()
            self.log_ai_interaction(prompt, content, True)
            return content
                
        except Exception as e:
            error_msg = f"AI模型调用失败: {str(e)}"
            print(error_msg)
            print(traceback.format_exc())
            self.log_ai_interaction(prompt, None, False, error_msg + '\n' + traceback.format_exc())
            print("请检查OLLAMA服务配置和连接")
            return None
    
    def get_valid_sql(self, question: str) -> (bool, str):
        """
        根据自然语言问题，多轮尝试生成并验证可执行的SQL语句。
        参数：
            question (str): 自然语言问题。
        返回：
            (bool, str): (是否成功, 最终SQL或失败原因)
        """
        max_retry = int(os.getenv('SQL_RETRY', 3))  # 可在.env中配置SQL_RETRY，默认3
        retry_count = 0
        prompt = self._generate_query_prompt(question)
        sql_query = self._get_ai_response(prompt)
        last_error = ''
        while retry_count < max_retry:
            valid, sql_checked = self.db.validate_and_fix_sql(sql_query)
            if valid:
                return True, sql_checked
            else:
                retry_count += 1
                last_error = sql_checked
                print(f"第{retry_count}次SQL验证失败，原因: {last_error}")
                prompt = self._generate_query_prompt(question) + f"\n上一次生成的SQL: {sql_query}\n验证失败原因: {last_error}\n请根据上述反馈重新生成正确的SQL语句。"
                sql_query = self._get_ai_response(prompt)
                if not sql_query:
                    break
        return False, last_error
    
    def query(self, question: str) -> Dict[str, Any]:
        """
        处理自然语言查询，自动生成 SQL 并执行。
        参数：
            question (str): 自然语言问题。
        返回：
            dict: 查询结果，包含 SQL、执行结果或错误信息。
        """
        success, sql_or_err = self.get_valid_sql(
            question,
            self._generate_query_prompt,
            self._get_ai_response,
            self.db.validate_and_fix_sql
        )
        if success:
            print(f"处理后的SQL: {sql_or_err}")
            try:
                results = self.db.execute_query(sql_or_err)
                return {
                    'success': True,
                    'sql': sql_or_err,
                    'results': results
                }
            except Exception as e:
                return {
                    'success': False,
                    'error': f'SQL执行失败: {str(e)}',
                    'traceback': traceback.format_exc()
                }
        else:
            return {
                'success': False,
                'error': f'SQL验证失败，重试后仍未通过: {sql_or_err}',
                'traceback': traceback.format_exc()
            }
